import mindspore as ms
import mindspore.nn as nn
import random
import numpy as np
from pytest import skip
try:
    from core import *
    from modules import *
except ImportError:
    from .core import *
    from .modules import *
    from .spectral_norm import SpectralNorm

class GSkip(nn.Cell):
    def __init__(self, skip_type, size, skip_init, skip_dropout=0, 
                 merge_mode='sum', kwidth=11, bias=True):
        # skip_init only applies to alpha skips
        super().__init__()
        self.merge_mode = merge_mode
        if skip_type == 'alpha' or skip_type == 'constant':
            if skip_init == 'zero':
                alpha_ = ms.numpy.zeros(size)
            elif skip_init == 'randn':
                alpha_ = ms.numpy.randn(size)
            elif skip_init == 'one':
                alpha_ = ms.numpy.ones(size)
            else:
                raise TypeError('Unrecognized alpha init scheme: ', 
                                skip_init)
            if skip_type == 'alpha':
                self.skip_k = ms.Parameter(alpha_.view(1, -1, 1))
            else:
                # constan, not learnable
                self.skip_k = ms.Parameter(alpha_.view(1, -1, 1))
                self.skip_k.requires_grad = False
        elif skip_type == 'conv':
            if kwidth > 1:
                pad = kwidth // 2
            else:
                pad = 0
            self.skip_k = nn.Conv1d(size, size, kwidth, stride=1, pad_mode= 'pad',padding=pad, has_bias=bias)
        else:
            raise TypeError('Unrecognized GSkip scheme: ', skip_type)
        self.skip_type = skip_type
        if skip_dropout > 0:
            self.skip_dropout = nn.Dropout(skip_dropout)

    def __repr__(self):
        if self.skip_type == 'alpha':
            return self.cls_name + '(Alpha(1))'
        elif self.skip_type == 'constant':
            return self.cls_name + '(Constant(1))'
        else:
            return super().__repr__()
    
    def construct(self, hj, hi):
        if self.skip_type == 'conv':
            sk_h = self.skip_k(hj)
        else:
            # skip_k = self.skip_k.repeat(hj.shape[0], 1, hj.shape[2])
            skip_k = ms.numpy.tile(self.skip_k, (hj.shape[0], 1, hj.shape[2]))
            sk_h = skip_k * hj
        if hasattr(self, 'skip_dropout'):
            sk_h = self.skip_dropout(sk_h)
        if self.merge_mode == 'sum':
            # merge with input hi on current layer
            return sk_h + hi
        elif self.merge_mode == 'concat':
            return ms.ops.concat((hi, sk_h), axis=1)
        else:
            raise TypeError('Unrecognized skip merge mode: ', self.merge_mode)

class Generator(Model):

    def __init__(self, ninputs, # 1
                 fmaps, # [64, 128, 256, 512, 1024]
                 kwidth, # 31
                 poolings, # [4, 4, 4, 4, 4]
                 dec_fmaps=None,
                 dec_kwidth=None,
                 dec_poolings=None,
                 z_dim=None, # 1024
                 no_z=False,
                 skip=True,
                 bias=False,
                 skip_init='one',
                 skip_dropout=0,
                 skip_type='alpha',
                 norm_type=None,
                 skip_merge='sum', # concat
                 skip_kwidth=11,
                 name='Generator'):
        super().__init__(name=name)
        self.skip = skip
        self.bias = bias
        self.no_z = no_z
        self.z_dim = z_dim
        # if not self.no_z:
        #     self.z = ms.Parameter(ms.Tensor(1.0, ms.float32), name="z")
        self.enc_blocks = []
        # 参数校验
        assert isinstance(fmaps, list), type(fmaps)
        assert isinstance(poolings, list), type(poolings)
        if isinstance(kwidth, int): 
            kwidth = [kwidth] * len(fmaps) # 扩展 kwidth 为list
        assert isinstance(kwidth, list), type(kwidth)

        skips = {}
        ninp = ninputs
        for pi, (fmap, pool, kw) in enumerate(zip(fmaps, poolings, kwidth),
                                              start=1):
            if skip and pi < len(fmaps):
                # Make a skip connection for all but last hidden layer
                gskip = GSkip(skip_type, fmap,
                              skip_init,
                              skip_dropout,
                              merge_mode=skip_merge,
                              kwidth=skip_kwidth,
                              bias=bias)
                l_i = pi - 1
                skips[l_i] = {'alpha':gskip}
                setattr(self, 'alpha_{}'.format(l_i), skips[l_i]['alpha']) # 向self类设置属性
            enc_block = GConv1DBlock( # 创建卷积层
                ninp, fmap, kw, stride=pool, bias=bias,
                norm_type=norm_type
            )
            self.enc_blocks.append(enc_block)
            ninp = fmap
        self.enc_blocks = nn.CellList(self.enc_blocks)

        self.skips = skips
        if not no_z and z_dim is None:
            z_dim = fmaps[-1]
        if not no_z:
            ninp += z_dim
        # Ensure we have fmaps, poolings and kwidth ready to decode
        if dec_fmaps is None: # 从encoder倒推decoder
            dec_fmaps = fmaps[::-1][1:] + [1]
        else:
            assert isinstance(dec_fmaps, list), type(dec_fmaps)
        if dec_poolings is None: # 从encoder倒推decoder
            dec_poolings = poolings[:]
        else:
            assert isinstance(dec_poolings, list), type(dec_poolings)
        self.dec_poolings = dec_poolings
        if dec_kwidth is None:
            dec_kwidth = kwidth[:]
        else:
            if isinstance(dec_kwidth, int): 
                dec_kwidth = [dec_kwidth] * len(dec_fmaps)
        assert isinstance(dec_kwidth, list), type(dec_kwidth)
        # Build the decoder
        self.dec_blocks = []
        for pi, (fmap, pool, kw) in enumerate(zip(dec_fmaps, dec_poolings, 
                                                  dec_kwidth),
                                              start=1):
            if skip and pi > 1 and pool > 1:
                if skip_merge == 'concat':
                    ninp *= 2

            if pi >= len(dec_fmaps):
                act = 'Tanh'
            else:
                act = None
            if pool > 1:
                dec_block = GDeconv1DBlock(
                    ninp, fmap, kw, stride=pool,
                    norm_type=norm_type, bias=bias,
                    act=act # None
                )
            else:
                dec_block = GConv1DBlock(
                    ninp, fmap, kw, stride=1, 
                    bias=bias,
                    norm_type=norm_type
                )
            self.dec_blocks.append(dec_block)
            ninp = fmap
        self.dec_blocks = nn.CellList(self.dec_blocks)

    def construct(self, x, z=None, ret_hid=False):
        hall = {}
        hi = x
        skips = self.skips
        for l_i, enc_layer in enumerate(self.enc_blocks):
            
            # print(f"x shape before enc: {hi.shape}")
            hi, linear_hi = enc_layer(hi, True)
            # print(f"x shape after enc: {hi.shape}")

            if self.skip and l_i < (len(self.enc_blocks) - 1):
                skips[l_i]['tensor'] = linear_hi
            if ret_hid:
                hall['enc_{}'.format(l_i)] = hi
        if not self.no_z:
            if z is None:
                # make z
                z = ms.numpy.randn(hi.shape[0], self.z_dim, *hi.shape[2:])
            if len(z.shape) != len(hi.shape):
                raise ValueError('len(z.size) {} != len(hi.size) {}'
                                 ''.format(len(z.shape), len(hi.shape)))
            if not hasattr(self, 'z'):
                self.z = z
            hi = ms.ops.concat((z, hi), axis=1)
            if ret_hid:
                hall['enc_zc'] = hi
        else:
            z = None
        enc_layer_idx = len(self.enc_blocks) - 1
        for l_i, dec_layer in enumerate(self.dec_blocks):
            if self.skip and enc_layer_idx in self.skips and self.dec_poolings[l_i] > 1:
                skip_conn = skips[enc_layer_idx] # 取出这一skip层
                hi = skip_conn['alpha'](skip_conn['tensor'], hi)
            hi = dec_layer(hi)
            enc_layer_idx -= 1
            if ret_hid:
                hall['dec_{}'.format(l_i)] = hi
        if ret_hid:
            return hi, hall
        else:
            return hi
        
if __name__ == '__main__':
    import os
    os.environ['CUDA_VISIBLE_DEVICES'] = '7'
    ms.set_context(mode=ms.PYNATIVE_MODE)
    G = Generator(1, [64, 128, 256, 512, 1024],
                  kwidth=31,
                  poolings=[4, 4, 4, 4, 4], no_z=True)
    # print(G)
    # ms.save_checkpoint(G,"ms_generator.ckpt")
    print('G num params: ', G.get_n_params())
    x = ms.numpy.ones((8, 1, 16384))
    print(x.shape)
    y = G(x)
    print(y)
    print(y.shape)